package solvers.linearsupport;

import globals.ExperimentTracker;
import globals.StopWatch;
import globals.Tuple;

import java.util.ArrayList;

import problems.DiscreteActions;
import problems.DiscreteObservations;
import problems.Mopomdp;
import solvers.vi.ValueIteratorPO;
import value.LazyScalarizedVector;
import value.ValueAtBeliefVector;

public class VanillaFullVISolver<S, A, O> implements SingleSolver<A> {

	public Mopomdp<S,A,O> problem;
	private double precision;
	private double[] initialBelief;
	private int cnt = 0;
	
	private boolean restart;
	private int restartCnt;
	private ArrayList<Tuple<Long,ValueAtBeliefVector<A>>> restartVectors;
	
	private static final boolean LOG = true;
	
	public VanillaFullVISolver(Mopomdp<S,A,O> m, double[] iBelief, double preciz){
		this.problem = m;
		this.precision = preciz;
		this.initialBelief = iBelief;
		this.cnt = 0;
		this.restart=false;
		this.restartCnt = 0;
		this.restartVectors = new ArrayList<Tuple<Long,ValueAtBeliefVector<A>>>();
	}
	
	public void setRestart(){
		this.restart = true;
		this.restartCnt = 0;
		this.restartVectors = new ArrayList<Tuple<Long,ValueAtBeliefVector<A>>>();
		
		double[] v1 = {13.3,-27.000000000000004};
		double[] v2 = {-8.717826075749288,-0.1500753498928722};
		Tuple<Long,ValueAtBeliefVector<A>> vec1 = new Tuple<Long,ValueAtBeliefVector<A>>((long) 49,new ValueAtBeliefVector<A>(v1));
		Tuple<Long,ValueAtBeliefVector<A>> vec2 = new Tuple<Long,ValueAtBeliefVector<A>>((long) 2543225,new ValueAtBeliefVector<A>(v2));
		this.restartVectors.add(vec1);
		this.restartVectors.add(vec2);
	}
	
	@Override
	public ValueAtBeliefVector<A> solveForW(CornerPoint<A> cp) {
		return this.solveForW(cp.weightVector);
	}

	@SuppressWarnings("unchecked")
	@Override
	public ValueAtBeliefVector<A> solveForW(double[] w) {
		
		ArrayList<LazyScalarizedVector<A>> lb = new ArrayList<LazyScalarizedVector<A>>();
		LazyScalarizedVector<A> lsv = new LazyScalarizedVector<A>(this.problem.misereMatrix(), w);
		lb.add(lsv);
		ValueIteratorPO<LazyScalarizedVector<A>,S,A,O> vi = 
				new ValueIteratorPO<LazyScalarizedVector<A>,S,A,O>(this.problem, lb);
		try{
			DiscreteActions<A,S> dacts   = (DiscreteActions<A,S>) problem.getActionSet();
			DiscreteObservations<O> dobs = (DiscreteObservations<O>) problem.getObservationSet();
			vi.valueIteration(dacts, dobs, problem.getStateSet(), this.precision);
			ValueAtBeliefVector<A> vb;
			if(this.restart && this.restartCnt<this.restartVectors.size()){
				vb = this.restartVectors.get(restartCnt).y;
				StopWatch sw = StopWatch.getGlobalStopWatch();
				sw.addTime(this.restartVectors.get(restartCnt).x);
				restartCnt++;
				cnt++;
			} else {
				vb = new ValueAtBeliefVector<A>(w.length, this.initialBelief, vi.alphas);
				if(LOG){
					StopWatch sw = StopWatch.getGlobalStopWatch();
					sw.stop();
					ExperimentTracker et = ExperimentTracker.getGlobalTracker(null);
					ArrayList<ValueAtBeliefVector<A>> vbs = new ArrayList<ValueAtBeliefVector<A>>(1);
					vbs.add(vb);
					et.writeVectorSet(""+(++cnt)+"_", vbs);
					Thread.sleep(100);
					sw.start();
				}
			}
			return vb;
		} catch(Exception e) {
			System.err.println("Trying to call discrete POMDP method for non-discrete problem.");
			System.err.println("Quitting...");
			System.exit(77);
			return null;
		}
		
	}
}
